package socket

import (
	"context"
	"errors"
	"github.com/chicken-team-outside/chicken_transmission/api_messages"
	"github.com/golang/protobuf/proto"
	"go.uber.org/zap"
	"slices"
	"sync"
)

var errSetClientListSuccess = errors.New("set client list success")

type clientCtxEntry struct {
	ctx          context.Context
	cancel       context.CancelCauseFunc
	info         *api_messages.TransportInfo
	onlineStatus *api_messages.TransportOnlineStatus
	lock         sync.Mutex
}

func newClientCtxEntry(ctx context.Context, info *api_messages.ClientStatusInfo) *clientCtxEntry {
	ctxEntry := &clientCtxEntry{
		info:         info.Info,
		onlineStatus: info.OnlineStatus,
	}
	ctxEntry.ctx, ctxEntry.cancel = context.WithCancelCause(ctx)
	if !info.OnlineStatus.Online {
		ctxEntry.cancel(ErrAckOffline)
	}
	return ctxEntry
}

func (c *WebSocketClient) setClientList(ctx context.Context, header *api_messages.MessageHeader, data []byte) {
	status := new(api_messages.ListClientReply)
	err := proto.Unmarshal(data, status)
	if err != nil {
		c.connCtxDone(err)
		return
	}
	newMap := make(map[uint64]*clientCtxEntry, len(status.Info))
	for _, info := range status.Info {
		newMap[info.Info.Id] = newClientCtxEntry(c.GetConn().GetCtx(), info)
	}
	c.clientCtxLock.Lock()
	c.clientCtx = newMap
	c.clientCtxLock.Unlock()
	c.connCtxDone(errSetClientListSuccess)
	c.notifyClientStatusChange(header.Type, status.Info)
}

func unmarshalClientStatusInfo(data []byte) *api_messages.ClientStatusInfo {
	info := new(api_messages.ClientStatusInfo)
	err := proto.Unmarshal(data, info)
	if err != nil {
		zap.L().Warn("failed to unmarshal client status info", zap.Error(err))
		return nil
	}
	return info
}

func unmarshalTransportOnlineStatus(data []byte) *api_messages.TransportOnlineStatus {
	status := new(api_messages.TransportOnlineStatus)
	err := proto.Unmarshal(data, status)
	if err != nil {
		zap.L().Warn("failed to unmarshal transport online status", zap.Error(err))
		return nil
	}
	return status
}

func (c *WebSocketClient) notifyClientStatusChange(msgType api_messages.MessageType, data interface{}) {
	callback := c.onStatusChange
	if callback == nil {
		return
	}
	callback(msgType, data)
}

func (c *WebSocketClient) getClientCtxEntry(id uint64) *clientCtxEntry {
	c.clientCtxLock.RLock()
	defer c.clientCtxLock.RUnlock()
	value, ok := c.clientCtx[id]
	if ok {
		return value
	}
	return nil
}

func (c *WebSocketClient) getClientCtx(clientId uint64) context.Context {
	if clientId == 0 {
		return c.GetCurConContext()
	}
	ctxEntry := c.getClientCtxEntry(clientId)
	if ctxEntry == nil {
		ctx, cancel := context.WithCancelCause(context.Background())
		cancel(ErrAckOffline)
		return ctx
	} else {
		return ctxEntry.ctx
	}
}

func (c *WebSocketClient) removeClientCtx(id uint64) {
	clientCtx := c.getClientCtxEntry(id)
	if clientCtx != nil {
		c.clientCtxLock.Lock()
		delete(c.clientCtx, id)
		c.clientCtxLock.Unlock()
		clientCtx.cancel(ErrAckOffline)
	}
}

func (c *WebSocketClient) onAddClient(ctx context.Context, header *api_messages.MessageHeader, data []byte) {
	info := unmarshalClientStatusInfo(data)
	if info == nil {
		return
	}
	c.clientCtxLock.Lock()
	c.clientCtx[info.Info.Id] = newClientCtxEntry(c.GetConn().GetCtx(), info)
	c.clientCtxLock.Unlock()
	c.notifyClientStatusChange(header.Type, info)
}

func (c *WebSocketClient) onRemoveClient(ctx context.Context, header *api_messages.MessageHeader, data []byte) {
	info := unmarshalTransportOnlineStatus(data)
	if info == nil {
		return
	}
	c.removeClientCtx(info.Id)
	c.notifyClientStatusChange(header.Type, info.Id)
}

func (c *WebSocketClient) onUpdateClientInfo(ctx context.Context, header *api_messages.MessageHeader, data []byte) {
	info := unmarshalClientStatusInfo(data)
	if info == nil {
		return
	}
	clientCtx := c.getClientCtxEntry(info.Info.Id)
	if clientCtx == nil {
		return
	}
	clientCtx.lock.Lock()
	defer clientCtx.lock.Unlock()
	hasChange := false
	if clientCtx.info.UpdateTime < info.Info.UpdateTime {
		clientCtx.info = info.Info
		hasChange = true
	}
	if clientCtx.onlineStatus.UpdateTime < info.OnlineStatus.UpdateTime {
		clientCtx.onlineStatus = info.OnlineStatus
		info.OnlineStatus.Online = clientCtx.ctx.Err() == nil
		hasChange = true
	}
	if hasChange {
		c.notifyClientStatusChange(header.Type, info)
	}
}

func (c *WebSocketClient) onClientStatusChange(ctx context.Context, header *api_messages.MessageHeader, data []byte) {
	info := unmarshalTransportOnlineStatus(data)
	if info == nil {
		return
	}
	clientCtx := c.getClientCtxEntry(info.Id)
	if clientCtx != nil {
		clientCtx.lock.Lock()
		if clientCtx.onlineStatus.UpdateTime >= info.UpdateTime {
			clientCtx.lock.Unlock()
			return
		}
		clientCtx.onlineStatus = info
		needNotify := false
		if info.Online {
			if clientCtx.ctx.Err() != nil {
				clientCtx.ctx, clientCtx.cancel = context.WithCancelCause(context.Background())
				needNotify = true
			}
		} else {
			if clientCtx.ctx.Err() == nil {
				clientCtx.cancel(ErrAckOffline)
				needNotify = true
			}
		}
		clientCtx.lock.Unlock()
		if needNotify {
			c.notifyClientStatusChange(header.Type, info)
		}
	}
}

func (c *WebSocketClient) GetClients() (result []*api_messages.ClientStatusInfo) {
	c.clientCtxLock.RLock()
	defer c.clientCtxLock.RUnlock()
	result = make([]*api_messages.ClientStatusInfo, 0, len(c.clientCtx))
	for _, item := range c.clientCtx {
		result = append(result, &api_messages.ClientStatusInfo{
			Info:         item.info,
			OnlineStatus: item.onlineStatus,
		})
	}
	slices.SortFunc(result, func(a, b *api_messages.ClientStatusInfo) int {
		if a.Info.Id > b.Info.Id {
			return 1
		}
		if a.Info.Id < b.Info.Id {
			return -1
		}
		return 0
	})
	return
}
